
import cv2
import torch
import torchvision
from torchvision.transforms import functional as F
from torchvision.ops import nms
from trainer import LyPoseDetector


device="cuda"
ckpt_path = "/home/disk0/hyq/JSPT/ly_pose2/ckpt/epochepoch=032-val_apval_ap=0.39.ckpt"
img_path = '/home/disk0/hyq/JSPT/data/1262/ZDSliedown20220720_V1_train_building_1_000167.jpg'

class ToTensor(object):
    def __call__(self, image, target):
        image = F.to_tensor(image)
        return image, target
data_trainsform = ToTensor()

model = LyPoseDetector()
model.load_from_checkpoint(ckpt_path)
img_original = cv2.imread(str(img_path))
img_original = cv2.cvtColor(img_original, cv2.COLOR_BGR2RGB)
data, _ = data_trainsform(img_original,target=None)

data = data.unsqueeze(0)
model.eval()
model.to(device)
data.to(device)
pred = model(data.to("cuda:0"))
pred = pred[0]

boxes, labels, keypoints, keypoints_scores = pred['boxes'], pred['labels'], pred['keypoints'], pred['keypoints_scores']
print(boxes.size())
print(labels)
print(keypoints.size())
print(keypoints_scores.size())


    